import os
import pymysql
import pandas as pd
from tqdm import tqdm


import sys
sys.path.append("/home/jie/.key")
# 本地密码存储文件
from sql_key import database, password


column_mapping = {
    '专利公开号': 'publication_number',
    '专利名称': 'patent_name',
    '专利类型': 'patent_type',
    '专利摘要': 'abstract',
    '申请人': 'applicant',
    '专利申请号': 'application_number',
    '申请日': 'application_date',
    '申请公布日': 'publication_date',
    '授权公布号': 'grant_publication_number',
    '授权公布日': 'grant_publication_date',
    '申请地址': 'application_address',
    '主权项': 'claims',
    '发明人': 'inventor',
    '分类号': 'classification_number',
    '主分类号': 'main_classification_number',
    '代理机构': 'agency',
    '分案原申请号': 'original_application_number',
    '优先权': 'priority',
    '国际申请': 'international_application',
    '国际公布': 'international_publication',
    '代理人': 'agent',
    '省份或国家代码': 'province_or_country_code',
    '法律状态': 'legal_status',
    '专利领域': 'patent_field',
    '专利学科': 'patent_subject',
    '多次公布': 'multiple_publications'
}


def trans2int(item):
    if pd.isna(item):
        return None
    try:
        return int(eval(item))
    except:
        return None


def parse_item(row, attr_name, args: list):
    
    assert isinstance(args, list)

    # 字符截断
    trunc_item = {
        "applicant" : 255,
        "application_address": 255,
        "inventor": 1024,
        "classification_number": 1024,
        "main_classification_number": 127,
        "agency" : 255,
        "original_application_number": 127,
        "priority" : 511,
        "international_application": 255,
        "international_publication": 255,
        "agent": 127,
        "patent_field" : 255,
        "patent_subject" : 255,
        "multiple_publications": 63,
    }

    ans = []
    for attr in attr_name:
        en_attr = column_mapping[attr]
        item = row.get(attr, None)
        
        if pd.isna(item):
            ans.append(None)
            continue
        
        if en_attr == "province_or_country_code":
            item = trans2int(item)
        # 异常字符捕获
        elif isinstance(item, str):
            item = item.strip()
            # 有 空串 和 -
            if len(item) in [0, 1]:
                item = None
        
        if en_attr in trunc_item.keys():
            if isinstance(item, str):
                item = item[:trunc_item[en_attr]]
                
        ans.append(item)
    ans += args
    
    return tuple(ans)


def insert_sql_by_csv(file):
    df = pd.read_csv(file, low_memory=False)
    
    province = os.path.basename(file).split('.')[0]
    # 连接到MySQL数据库
    connection = pymysql.connect(
        host="localhost",  # MySQL数据库的主机
        user="root",  # MySQL用户名
        password=password,  # MySQL密码
        database=database,  # 你要插入数据的数据库
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
    )
    
    BATCH_SIZE = 1000

    tmp_list = list(column_mapping.values()) + ["province"]
    sql = f"""
            INSERT INTO patents (
                {", ".join(tmp_list)}
            ) VALUES (
                {', '.join(['%s'] * len(tmp_list))}
            );
            """.strip()
    
    # 插入数据到MySQL
    try:
        with connection.cursor() as cursor:
            batch = []
            for _, row in tqdm(df.iterrows(), total=len(df)):
    
                batch.append(
                        parse_item(row, list(column_mapping.keys()), [province])
                    )

                # 当批次达到 BATCH_SIZE 时执行批量插入
                if len(batch) >= BATCH_SIZE:
                    cursor.executemany(sql, batch)
                    batch = []  # 清空批次
                    
            # 插入剩余的未满批次的数据
            if batch:
                cursor.executemany(sql, batch)

            # 提交事务
            connection.commit()
            
    except Exception as e:
        print(f"插入数据时出现错误: {e}")
        connection.rollback()

    finally:
        connection.close()


if __name__ == "__main__":    
    folder = "/media/jie/新加卷/pku_data/3571万专利申请全量数据1985-2022年/"
    print(f"文件总数: {len(os.listdir(folder))}")
    cnt = 0
    for file_name in os.listdir(folder):
        if file_name.endswith(".csv"):
            cnt += 1
            filename = os.path.join(folder, file_name)
            print(cnt, file_name)
            insert_sql_by_csv(filename)

    print("所有数据全部插入完成！")

    # mysql -h 127.0.0.1 -P 3306 -u root -p
    # nohup python build_patent_table.py > build_patent_table.log 2>&1 &






